package com.jonlandrum.collections.BinarySearchTree.AVLTree;

import com.jonlandrum.collections.BinarySearchTree.LinkedAdjustableTreeADT;

import java.util.NoSuchElementException;

public class LinkedAVLTree<T extends Comparable<T>> extends LinkedAdjustableTreeADT<T> implements AVLTree<T> {
    /**
     * Constructs an empty AVLTree.
     */
    public LinkedAVLTree() {
        super();
    }

    /**
     * Constructs a new AVLTree with the given node set as the root element.
     *
     * @param node The node to set as this tree's root element
     */
    public LinkedAVLTree(AVLNode<T> node) {
        super(node);
    }

    /**
     * Returns the root node of this tree.
     *
     * @return The root node of this tree
     */
    @Override
    public AVLNode<T> getRoot() {
        return (LinkedAVLNode<T>) super.getRoot();
    }

    /**
     * Inserts a new data element into this tree.
     *
     * @param data The data element to insert into this tree
     */
    @Override
    public AVLNode<T> insert(T data) {
        return this.insert(new LinkedAVLNode<>(data));
    }

    /**
     * Inserts a new node into this tree.
     *
     * @param node The node to insert into this tree
     */
    @Override
    public AVLNode<T> insert(AVLNode<T> node) {
        super.insert(node);
        node.setBalanceFactor();
        this.balance(node);
        return node;
    }

    /**
     * Deletes a data element from this tree.
     *
     * @param data The data element to be removed from this tree
     * @return The replacement of the data element removed from this tree
     */
    @Override
    public AVLNode<T> delete(T data) {
        return this.delete(new LinkedAVLNode<>(data));
    }

    /**
     * Deletes a node from this tree.
     *
     * @param node The node to be removed from this tree
     * @return The replacement of the node removed from this tree
     */
    @Override
    public AVLNode<T> delete(AVLNode<T> node) {
        LinkedAVLNode<T> target = (LinkedAVLNode<T>) search(node);
        LinkedAVLNode<T> replacement = (LinkedAVLNode<T>) target.getReplacement();
        super.delete(target);
        if (replacement.hasData()) {
            replacement.setBalanceFactor();
            this.balance(replacement);
        }
        return replacement;
    }

    /**
     * Searches for a specific data element in this tree.
     *
     * @param data The data element of the same value as the node searched for
     * @return The node searched for
     */
    @Override
    public AVLNode<T> search(T data) {
        return this.search(new LinkedAVLNode<>(data));
    }

    /**
     * Searches for a specific node in this tree.
     *
     * @param node The data element of the same value as the node searched for
     * @return The node searched for
     * @throws NoSuchElementException if the desired node is not in this tree
     */
    @Override
    public AVLNode<T> search(AVLNode<T> node) throws NoSuchElementException {
        LinkedAVLNode<T> result = (LinkedAVLNode<T>) super.search(node);
        if (!hasRoot() ||
                result == null ||
                result.compareTo(node) != 0) {
            throw new NoSuchElementException("No Such Element");
        }
        return result;
    }

    private void balance(AVLNode<T> node) {
        int balance = node.getBalanceFactor();
        while ((Math.abs(balance) > 1) || node.hasParent()) {
            if (balance < -1) {
                this.rotateRight(node);
            } else if (balance > 1) {
                this.rotateLeft(node);
            }
            if (node.hasParent()) {
                node = (LinkedAVLNode<T>) node.getParent();
            }
            balance = node.getBalanceFactor();
        }
    }
}
